import torch
import torch.nn as nn
import pytorch_model_zoo
from torch.autograd import Variable
import numpy as np

# 准备模型DPN107 for Class

USE = 'Rank'

CKPT = '/mnt/md1/Downloads/dpn107_80e_72.40p.pth.tar'

ckpt = torch.load(CKPT)
base_model = getattr(pytorch_model_zoo, 'dpn107')()
base_model.last_layer_name = 'classifier' #attention:last layer is a conv

if USE == 'Class':
    setattr(base_model, base_model.last_layer_name, nn.Conv2d(2688,201,kernel_size=(1,1),stride=(1,1)))
elif USE == 'Rank':
    setattr(base_model, base_model.last_layer_name, nn.Conv2d(2688,2,kernel_size=(1,1),stride=(1,1)))

ckpt_keys = list(ckpt['state_dict'].keys())
ckpt_keys = list(map(lambda x: '.'.join(x.split('.')[1:]),list(filter(lambda x: x.split('.')[0]=='base_model',
                                                                      list(map(lambda x: '.'.join(x.split('.')[1:]),ckpt_keys))))))

cnt = 0
init = 0
DDD = dict()
for key in base_model.state_dict():
    cnt += 1
    if key in ckpt_keys:
        init+=1
        DDD[key] = ckpt['state_dict']['module.base_model.'+key]
    else:
        print(key)
print(init,'/',cnt)

base_model.state_dict().update(DDD)


input_size = 224
input_mean =  [124. / 255, 117. / 255, 104. / 255]
input_std = [1 / (.0167 * 255)] * 3

I = Variable(torch.from_numpy(np.random.rand(10,3,224,224))).float()
O = base_model(I)

if USE=='Class':
    torch.save(base_model.state_dict(),'/mnt/md1/Experiments/DPN_Extruct_200_Test1/dpn107_state_dict.pkl')
    torch.save(base_model,'/mnt/md1/Experiments/DPN_Extruct_200_Test1/raw_dpn107_model.pkl')
elif USE=='Rank':
    torch.save(base_model.state_dict(),'/mnt/md1/Experiments/DPN_Extruct_200_Test1/dpn107_rank_state_dict.pkl')
    torch.save(base_model,'/mnt/md1/Experiments/DPN_Extruct_200_Test1/raw_dpn107_rank_model.pkl')

base_model = torch.load('/mnt/md1/Experiments/DPN_Extruct_200_Test1/raw_dpn107_model.pkl')


